from abc import ABC, abstractmethod
from typing import Dict, Any, Type, Optional, Callable
import time
import torch.nn as nn
from ..config import Config
from ..utils import setup_logger

class ModelCompressor(ABC):
    """Abstract base class for model compressors.
    
    Attributes:
        model: The neural network model to compress.
        config: Configuration parameters for compression.
        logger: Logger instance for tracking operations.
        metrics: Dictionary to store compression metrics.
    """
    
    _compressors: Dict[str, Type['ModelCompressor']] = {}  # Registry of available compressors
    
    def __init__(self, model: nn.Module, config: Config):
        """Initialize compressor with model and configuration.
        
        Args:
            model: PyTorch model to compress.
            config: Compression configuration parameters.
            
        Raises:
            TypeError: If model is not a nn.Module or config is not Config.
        """
        if not isinstance(model, nn.Module):
            raise TypeError(f"Expected nn.Module, got {type(model)}")
        if not isinstance(config, Config):
            raise TypeError(f"Expected Config, got {type(config)}")
            
        self.model = model
        self.config = config
        self.logger = setup_logger(self.__class__.__name__)
        self.metrics: Dict[str, Any] = {
            'compression_time': 0.0,
            'original_size': 0,
            'compressed_size': 0
        }
        
    @classmethod
    def register_compressor(cls, name: str) -> Callable[[Type['ModelCompressor']], Type['ModelCompressor']]:
        """Decorator to register a compressor class with validation.
        
        Args:
            name: Unique name for the compressor type.
            
        Returns:
            Decorator function.
            
        Raises:
            TypeError: If subclass is not a ModelCompressor.
            ValueError: If name is already registered.
            
        Example:
            @ModelCompressor.register_compressor('pruning')
            class PruningCompressor(ModelCompressor):
                ...
        """
        def decorator(subclass: Type['ModelCompressor']) -> Type['ModelCompressor']:
            if not issubclass(subclass, cls):
                raise TypeError(f"Compressor must subclass ModelCompressor, got {subclass}")
            if name in cls._compressors:
                raise ValueError(f"Compressor '{name}' already registered")
                
            cls._compressors[name] = subclass
            subclass.compressor_name = name  # Store name in class
            return subclass
        return decorator
    
    @classmethod
    def get_compressor(cls, name: str) -> Optional[Type['ModelCompressor']]:
        """Get registered compressor class by name.
        
        Args:
            name: Name of the registered compressor.
            
        Returns:
            Compressor class if found, None otherwise.
        """
        return cls._compressors.get(name)
    
    @classmethod
    def list_compressors(cls) -> Dict[str, str]:
        """Get all registered compressors with descriptions.
        
        Returns:
            Dictionary mapping compressor names to their docstring summaries.
        """
        return {
            name: (comp.__doc__ or '').split('\n')[0]
            for name, comp in cls._compressors.items()
        }
    
    @abstractmethod
    def compress(self) -> nn.Module:
        """Compress the model while tracking performance metrics.
        
        Returns:
            Compressed model.
            
        Note:
            Implementations should update self.metrics with relevant statistics.
        """
        pass
    
    @abstractmethod
    def evaluate(self, test_loader) -> Dict[str, Any]:
        """Evaluate compressed model on test data.
        
        Args:
            test_loader: Data loader for evaluation.
            
        Returns:
            Dictionary containing:
                - accuracy: Model accuracy
                - loss: Model loss
                - inference_time: Average inference time
                - memory_usage: Memory consumption
            
        Raises:
            RuntimeError: If evaluation fails.
        """
        pass
    
    def _apply_compression(self) -> nn.Module:
        """Internal method to apply compression with timing and logging.
        
        Returns:
            Compressed model.
            
        Raises:
            RuntimeError: If compression fails.
        """
        self.logger.info("Starting compression...")
        start_time = time.time()
        
        try:
            compressed_model = self.compress()
            compression_time = time.time() - start_time
            self.metrics['compression_time'] = compression_time
            self.logger.info(f"Compression completed in {compression_time:.2f}s")
            return compressed_model
        except Exception as e:
            self.logger.error(f"Compression failed: {str(e)}")
            raise RuntimeError(f"Compression failed: {str(e)}") from e
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get compression metrics.
        
        Returns:
            Dictionary of compression metrics including:
                - compression_time: Time taken for compression
                - original_size: Original model size (params)
                - compressed_size: Compressed model size (params)
                - compression_ratio: Size reduction ratio
        """
        metrics = self.metrics.copy()
        if 'original_size' in metrics and 'compressed_size' in metrics:
            metrics['compression_ratio'] = (
                metrics['original_size'] / metrics['compressed_size']
                if metrics['compressed_size'] > 0 else 0
            )
        return metrics

def create_compressor(name: str, model: nn.Module, config: Config) -> ModelCompressor:
    """Create and initialize a compressor instance.
    
    Args:
        name: Name of the registered compressor.
        model: PyTorch model to compress.
        config: Compression configuration parameters.
        
    Returns:
        Initialized compressor instance.
        
    Raises:
        ValueError: If compressor name is not registered.
        TypeError: If model or config types are invalid.
        RuntimeError: If compressor initialization fails.
    """
    compressor_cls = ModelCompressor.get_compressor(name)
    if compressor_cls is None:
        available = ModelCompressor.list_compressors()
        raise ValueError(
            f"Unknown compressor '{name}'. Available compressors: {available}"
        )
    
    if not isinstance(model, nn.Module):
        raise TypeError(f"Expected nn.Module for model, got {type(model)}")
    if not isinstance(config, Config):
        raise TypeError(f"Expected Config for config, got {type(config)}")
    
    try:
        return compressor_cls(model, config)
    except Exception as e:
        raise RuntimeError(
            f"Failed to initialize {name} compressor: {str(e)}"
        ) from e